
#include "color_quantization.h"
#include <string.h>
#include <cmath>

void MedianSegAnyColorNew(float *xyz_data, int width, int height, int seg_label, unsigned char *out_data, ColorRange *cr) {
	int img_width_step = CALCULATE_STEP_THREE(width);
	float x_channel_max = -1000;
    float x_channel_min = 1000;

    float y_channel_max = -1000;
    float y_channel_min = 1000;

    float z_channel_max = -1000;
    float z_channel_min = 1000;
	int x, y;
	float x_length, y_length, z_length;
	float seg_th;
	unsigned char tmp_label;
	float a, b, c;

	for (y = 0; y < height; y++) {
		for (x = 0; x < width; x++) {
			if (IMAGE_ELEM(out_data, unsigned char, width, x, y) == seg_label) {
				a = IMAGE_ELEM(xyz_data, float, img_width_step, x * 3 + 0, y);
				b = IMAGE_ELEM(xyz_data, float, img_width_step, x * 3 + 1, y);
				c = IMAGE_ELEM(xyz_data, float, img_width_step, x * 3 + 2, y);

				x_channel_max = x_channel_max > a ? x_channel_max : a;
				x_channel_min = x_channel_min < a ? x_channel_min : a;

				y_channel_max = y_channel_max > b ? y_channel_max : b;
				y_channel_min = y_channel_min < b ? y_channel_min : b;

				z_channel_max = z_channel_max > c ? z_channel_max : c;
				z_channel_min = z_channel_min < c ? z_channel_min : c;
			}
		}
	}
	
	x_length = (x_channel_max - x_channel_min - cr->x_min) / (cr->x_max - cr->x_min);
    y_length = (y_channel_max - y_channel_min - cr->y_min) / (cr->y_max - cr->y_min);
    z_length = (z_channel_max - z_channel_min - cr->z_min) / (cr->z_max - cr->z_min);

    if (z_length >= x_length && z_length >= y_length) {
        seg_th = (z_channel_max - z_channel_min) / 2 + z_channel_min;
		for (y = 0; y < height; y++) {
			for (x = 0; x < width; x++) {
				tmp_label = IMAGE_ELEM(out_data, unsigned char, width, x, y);
				if (tmp_label == seg_label) {
					if (IMAGE_ELEM(xyz_data, float, img_width_step, x * 3 + 2, y) >= seg_th) {
						IMAGE_ELEM(out_data, unsigned char, width, x, y) = tmp_label * 2 + 1 + CLUSTER_SIZE;
					} else {
						IMAGE_ELEM(out_data, unsigned char, width, x, y) = tmp_label * 2 + CLUSTER_SIZE;
					}
				}
			}
		}
    } else if (x_length >= z_length && x_length >= y_length) {
        seg_th = (x_channel_max - x_channel_min) / 2 + x_channel_min;
		for (y = 0; y < height; y++) {
			for (x = 0; x < width; x++) {
				tmp_label = IMAGE_ELEM(out_data, unsigned char, width, x, y);
				if (tmp_label == seg_label) {
					if (IMAGE_ELEM(xyz_data, float, img_width_step, x * 3 + 0, y) >= seg_th) {
						IMAGE_ELEM(out_data, unsigned char, width, x, y) = tmp_label * 2 + 1 + CLUSTER_SIZE;
					} else {
						IMAGE_ELEM(out_data, unsigned char, width, x, y) = tmp_label * 2 + CLUSTER_SIZE;
					}
				}
            }
        }
    } else {
		seg_th = (y_channel_max - y_channel_min) / 2 + y_channel_min;
		for (y = 0; y < height; y++) {
			for (x = 0; x < width; x++) {
				tmp_label = IMAGE_ELEM(out_data, unsigned char, width, x, y);
				if (tmp_label == seg_label) {
					if (IMAGE_ELEM(xyz_data, float, img_width_step, x * 3 + 1, y) >= seg_th) {
						IMAGE_ELEM(out_data, unsigned char, width, x, y) = tmp_label * 2 + 1 + CLUSTER_SIZE;
					} else {
						IMAGE_ELEM(out_data, unsigned char, width, x, y) = tmp_label * 2 + CLUSTER_SIZE;
					}
				}
			} 
		} 
	}
}

void MedianCutAnyColorNew(float *xyz_data, int width, int height, unsigned char *out_data, RGBCluster *rgbcluster, ColorRange *cr) {
	int iter_num = 0;
	int max_label;
	int i, x, y;
	int seg_label = 0;
	float xx = 0, yy = 0, zz = 0;
	int count = 0;
	RGBCluster cl;
	int img_width_step = CALCULATE_STEP_THREE(width);

	while (iter_num < 4) {
		max_label = (1 << iter_num);
		for (i = 0; i < max_label; i++) {
			MedianSegAnyColorNew(xyz_data, width, height, i, out_data, cr);
		}
		for (y = 0; y < height; y++) {
			for (x = 0; x < width; x++) {
				IMAGE_ELEM(out_data, unsigned char, width, x, y) -= CLUSTER_SIZE;
			}
		}
		iter_num++;
	} // end while
	
	while (seg_label < 16) {
		for (y = 0; y < height; y++) {
			for (x = 0; x < width; x++) {
				if (IMAGE_ELEM(out_data, unsigned char, width, x, y) == seg_label) {
					xx += IMAGE_ELEM(xyz_data, float, img_width_step, x * 3 + 0, y);
					yy += IMAGE_ELEM(xyz_data, float, img_width_step, x * 3 + 1, y);
					zz += IMAGE_ELEM(xyz_data, float, img_width_step, x * 3 + 2, y);
					count++;
				}
			}
		}
		
		if (count == 0) {
			cl.b = FLOAT_MAX;
			cl.g = FLOAT_MAX;
			cl.r = FLOAT_MAX;
		} else {
			cl.r = xx / count;
            cl.g = yy / count;
            cl.b = zz / count;
		}
		rgbcluster[seg_label].r = cl.r;
		rgbcluster[seg_label].g = cl.g;
		rgbcluster[seg_label].b = cl.b;

		seg_label++;
		xx = 0;
		yy = 0;
		zz = 0;
		count = 0;
	}//end while
}


void GetXYZClusterInfo(unsigned char *img_data, int width, int height, float *xyz_data, unsigned char *out_data, RGBCluster *rgbcluster) {
	int img_width_step = CALCULATE_STEP_THREE(width);
	int half_width = (width + 1) / 2;
	int half_height = (height + 1) / 2;
	int half_width_step = CALCULATE_STEP_THREE(half_width);
	int res_width_step = width;
	int r, g, b;
	float oa, ob, oc;
	int x, y;
	ColorRange cr;
	int sy, sx;
	int yy, xx;
	int count = 0;


	memset(out_data, 0, sizeof(unsigned char) * width * height);
	for (y = 0, sy = 0; y < height; y += 2, sy++) {
		for (x = 0, sx = 0; x < width; x += 2, sx++) {
			
			r = IMAGE_ELEM(img_data, unsigned char, img_width_step, x * 3 + 2, y);
			g = IMAGE_ELEM(img_data, unsigned char, img_width_step, x * 3 + 1, y);
			b = IMAGE_ELEM(img_data, unsigned char, img_width_step, x * 3 + 0, y);
			RGB2XYZ(r, g, b, &oa, &ob, &oc);
			IMAGE_ELEM(xyz_data, float, half_width_step, sx * 3 + 0, sy) = oa;
			IMAGE_ELEM(xyz_data, float, half_width_step, sx * 3 + 1, sy) = ob;
			IMAGE_ELEM(xyz_data, float, half_width_step, sx * 3 + 2, sy) = oc;
		}
	} //end outer for
	cr.x_max = 100.0f;
    cr.x_min = -100.0f;
    cr.y_max = 100.0f;
    cr.y_min = -100.0f;
    cr.z_max = 100.0f;
    cr.z_min = 0.0f;
	
	half_height = sy;
	half_width = sx;
	MedianCutAnyColorNew(xyz_data, half_width, half_height, out_data, rgbcluster, &cr);

	CreateXYZClusterMap(img_data, out_data, width, height, rgbcluster);
}

int FindNearestClusterNew(RGBCluster *rgbcluster, float oc_a, float oc_b, float oc_c) {
	float min_dist = FLOAT_MAX;
	int min_dist_index = -1;
	float a, b, c, dist;
	int i;
	for (i = 0; i < CLUSTER_SIZE; i++) {
		c = rgbcluster[i].b;
		b = rgbcluster[i].g;
		a = rgbcluster[i].r;
		dist = (a - oc_a) * (a - oc_a) + (b - oc_b) * (b - oc_b) + (c - oc_c) * (c - oc_c);
		if (dist < min_dist) {
			min_dist = dist;
			min_dist_index = i;
		}
	}//end for
	return min_dist_index;
}


void CreateXYZClusterMap(unsigned char *img_data, unsigned char *out_data, int width, int height, RGBCluster *rgbcluster) {
	float oa, ob, oc;
	int img_width_step = CALCULATE_STEP_THREE(width);
	int res_width_step = width;
	int x, y, i;
	int r, g, b;

	for (y = 0; y < height; y++) {
		for (x = 0; x < width; x++) {
			r = IMAGE_ELEM(img_data, unsigned char, img_width_step, x * 3 + 2, y);
			g = IMAGE_ELEM(img_data, unsigned char, img_width_step, x * 3 + 1, y);
			b = IMAGE_ELEM(img_data, unsigned char, img_width_step, x * 3 + 0, y);
			RGB2XYZ(r, g, b, &oa, &ob, &oc);
			i = FindNearestClusterNew(rgbcluster, oa, ob, oc);
			IMAGE_ELEM(out_data, unsigned char, res_width_step, x, y) = (unsigned char)i;
		}
	}
}

void RGB2XYZ(int r, int g, int b, float *x, float *y, float *z) {
	int max_v, min_v, diff;
	float hf = 0;
	float radian;
	int h, s, v;
	max_v = MAX(r, MAX(g, b));
	min_v = MIN(r, MIN(g, b));
	diff = max_v - min_v;
	//hue
	if (diff == 0) {
		hf = 0.0f;
	} else {
		if (r == max_v) {
			hf = (g - b) / (float)diff;
		}
		if (g == max_v) {
			hf = 2 + (b - r) / (float)diff;
		}
		if (b == max_v) {
			hf = 4 + (r - g) / (float)diff;
		}
		hf *= 60.0f;
		if (hf < 0) {
			hf += 360.0f;
		}
	}
	h = (int)(hf + 0.5);
	//saturation
	if (max_v == 0) {
		s = 0;
	} else {
		s = diff * 100 / max_v;
	}
	//value
	v = 100 * max_v / 255;
	if (v < 10) {
		h = 0, s = 0;
	}
	radian = (float)h * PI / 180.0;
	*x = s * cos(radian);
	*y = s * sin(radian);
	*z = v;
}




int FindNearestClusterNewByOneD(RGBCluster *rgbcluster, float oc_a, float oc_b, float oc_c) {
	float min_dist = FLOAT_MAX;
	int min_dist_index = -1;
	float a, b, c, dist;
	int i;
	for (i = 0; i < CLUSTER_SIZE; i++) {
		c = rgbcluster[i].b;
		b = rgbcluster[i].g;
		a = rgbcluster[i].r;
		dist = (a - oc_a) * (a - oc_a) + (b - oc_b) * (b - oc_b) + (c - oc_c) * (c - oc_c);
		if (dist < min_dist) {
			min_dist = dist;
			min_dist_index = i;
		}
	}//end for
	return min_dist_index;
}
void CreateXYZClusterMapByOneD(float *xyz_data, unsigned char *out_data, int width, int height, RGBCluster *rgbcluster) {
	float a, b, c;
	int img_width_step = CALCULATE_STEP_THREE(width);
	int res_width_step = width;
	int x, y, i;

	for (y = 0; y < height; y++) {
		for (x = 0; x < width; x++) {
			a = IMAGE_ELEM(xyz_data, float, img_width_step, x * 3 + 0, y);
			b = IMAGE_ELEM(xyz_data, float, img_width_step, x * 3 + 1, y);
			c = IMAGE_ELEM(xyz_data, float, img_width_step, x * 3 + 2, y);
			i = FindNearestClusterNewByOneD(rgbcluster, a, b, c);
			IMAGE_ELEM(out_data, unsigned char, res_width_step, x, y) = (unsigned char)i;
		}
	}
}
void GetXYZClusterInfoNew(unsigned char *img_data, int width, int height, float *xyz_data, unsigned char *out_data, RGBCluster *rgbcluster) {
	int img_width_step = CALCULATE_STEP_THREE(width);
	int res_width_step = width;
	int r, g, b;
	float oa, ob, oc;
	int x, y;
	ColorRange cr;

	for (y = 0; y < height; y++) {
		for (x = 0; x < width; x++) {
			
			r = IMAGE_ELEM(img_data, unsigned char, img_width_step, x * 3+ 2, y);
			g = IMAGE_ELEM(img_data, unsigned char, img_width_step, x * 3+ 1, y);
			b = IMAGE_ELEM(img_data, unsigned char, img_width_step, x * 3+ 0, y);
			RGB2XYZ(r, g, b, &oa, &ob, &oc);
			IMAGE_ELEM(out_data, unsigned char, res_width_step, x, y) = 0;
			IMAGE_ELEM(xyz_data, float, img_width_step, x * 3 + 0, y) = oa;
			IMAGE_ELEM(xyz_data, float, img_width_step, x * 3 + 1, y) = ob;
			IMAGE_ELEM(xyz_data, float, img_width_step, x * 3 + 2, y) = oc;
		}
	} //end outer for
	cr.x_max = 100.0f;
    cr.x_min = -100.0f;
    cr.y_max = 100.0f;
    cr.y_min = -100.0f;
    cr.z_max = 100.0f;
    cr.z_min = 0.0f;
	
	MedianCutAnyColorNew(xyz_data, width, height, out_data, rgbcluster, &cr);
	CreateXYZClusterMapByOneD(xyz_data, out_data, width, height, rgbcluster);
}